from towhee import ops
from pymilvus import connections, Collection
from flask import Flask, request, jsonify

config = {'towhee': {'model': 'resnet50', 'device': None}, 'milvus': {'host': '127.0.0.1', 'port': '19530',
    'dim': 2048, 'table_name': 'image_search', 'index_type': 'IVF_FLAT', 'metric_type': 'L2'}}

app = Flask(__name__)


def img_query(img_src, config):
    # Connect to Milvus service
    connections.connect(host=config["milvus"]["host"], port=config["milvus"]["port"])
    data_table = Collection(config["milvus"]["table_name"])
    data_table.load(replica_number=1)

    img = ops.image_decode()(img_src)
    vec = ops.image_embedding.timm(model_name=config["towhee"]["model"], device=config["towhee"]["device"])(img)

    search_param = {"metric_type": config["milvus"]["metric_type"], "offset": 0}
    result = data_table.search([vec], anns_field='vec', param=search_param, limit=20, output_fields=["id", "name", "uniq_no"])
    return result


@app.route("/api", methods=["POST"])
def api_search():
    img_src = request.form.get("s")
    result = img_query(img_src, config)
    return jsonify({"code": 0, "datas": result})


if __name__ == "__main__":
    app.run(port=9000)
